{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "ba282287",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LlamaForCausalLM(\n",
       "  (model): LlamaModel(\n",
       "    (embed_tokens): Embedding(32000, 256)\n",
       "    (layers): ModuleList(\n",
       "      (0-3): 4 x LlamaDecoderLayer(\n",
       "        (self_attn): LlamaAttention(\n",
       "          (q_proj): Linear(in_features=256, out_features=256, bias=False)\n",
       "          (k_proj): Linear(in_features=256, out_features=128, bias=False)\n",
       "          (v_proj): Linear(in_features=256, out_features=128, bias=False)\n",
       "          (o_proj): Linear(in_features=256, out_features=256, bias=False)\n",
       "          (rotary_emb): LlamaRotaryEmbedding()\n",
       "        )\n",
       "        (mlp): LlamaMLP(\n",
       "          (gate_proj): Linear(in_features=256, out_features=768, bias=False)\n",
       "          (up_proj): Linear(in_features=256, out_features=768, bias=False)\n",
       "          (down_proj): Linear(in_features=768, out_features=256, bias=False)\n",
       "          (act_fn): SiLU()\n",
       "        )\n",
       "        (input_layernorm): LlamaRMSNorm()\n",
       "        (post_attention_layernorm): LlamaRMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (norm): LlamaRMSNorm()\n",
       "  )\n",
       "  (lm_head): Linear(in_features=256, out_features=32000, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoModelForCausalLM\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "model_file = r'C:\\Users\\luyingli\\Desktop\\project\\llama\\saves\\checkpoint-177000'\n",
    "model = AutoModelForCausalLM.from_pretrained(model_file).to(device)\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "56195cfe",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BertTokenizerFast(name_or_path='C:\\Users\\luyingli\\Desktop\\project\\llama\\saves\\checkpoint-177000', vocab_size=21128, model_max_length=1000000000000000019884624838656, is_fast=True, padding_side='right', truncation_side='right', special_tokens={'unk_token': '[UNK]', 'sep_token': '[SEP]', 'pad_token': '[PAD]', 'cls_token': '[CLS]', 'mask_token': '[MASK]'}, clean_up_tokenization_spaces=True),  added_tokens_decoder={\n",
       "\t0: AddedToken(\"[PAD]\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
       "\t100: AddedToken(\"[UNK]\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
       "\t101: AddedToken(\"[CLS]\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
       "\t102: AddedToken(\"[SEP]\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
       "\t103: AddedToken(\"[MASK]\", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),\n",
       "}"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "# tokenizer = AutoTokenizer.from_pretrained('./chinese_llama')\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_file, trust_remote_code=True)\n",
    "\n",
    "tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "e69ce623",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "类 型 # 上 衣 * 版 型 # 宽 松 * 颜 色 # 粉 红 色 * 图 案 # 字 母 * 图 案 # 文 字 * 图 案 # 线 条 * 衣 样 式 # 卫 衣 * 衣 款 式 # 不 规 则 身 身 计 线 ， 色 ， ， 衣 身 。 母 的 ， ， 色 的 字 粉 ， 的 的 部 色 ， 计 的 ， 母 色 的 。 母 前 的 的 ， ， 字 色 的 的 ， 衣 的 设 ， 领 条 ， 有 色 的 。 了 缀 ， 的 ， 身 饰 饰 ， ， ， ， 色 ， ， 的 ， 衣 摆 的 。 口 。 衣 设 。 色 ， 口 的 ， 饰 母 前 ， 身 计 。 感 字 ， 字 计 ， 的 ， 。 一 的 ， 的 设 。 母 的 了 。 。 ， ， 色 ， 母 ， 衣 感 ， 显 。 色 的 ， 的 ， 的 的 的 ， 的 的 色 少 ， 单 的 的 色 少 有 。 色 感 。 的 母 的 的 的 体 的 母 ， 色 。 身 色 ， 型 的 身 的 身 力 。 卫 饰 计 领 设 时 ， 。 ， 身 ， 的 出 计 。 ， 衫 ， 长 ， 性 了 ， 条 色 ， ， 计 身 饰 的 显 ， 单 ， 计 ， 的 母 颈 领 ， 母 。 的 前 计 色 的 色 个 感 母 ， 字 色 母 了 ， 。 ， 衣 的 的 有 色 气 ， 字 。 得 的 ， ， ， ， 的 ， 衣 了 领 。 ， 点 。 色 ， 。 有 性 显 ， 母 ， 的 的 感 的 ， 母 口 ， ， ， 的 母 。 时 ， ， ， ， 饰 ， 。 的 的 衫 。 ， 母 ， 的 衫 的 的 。 是 有 的 ， 的 ， 。 了 的 的 气 ， ， 小 了\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "# 设置环境变量\n",
    "os.environ['CUDA_LAUNCH_BLOCKING'] = '1'\n",
    "\n",
    "\n",
    "def inference(\n",
    "    model: AutoModelForCausalLM,    \n",
    "    tokenizer: AutoTokenizer,    \n",
    "    input_text: str = '那咋啦 ',    \n",
    "    max_new_tokens: int = 16\n",
    "):\n",
    "    inputs = tokenizer(input_text, return_tensors='pt').to(device)  \n",
    "    if 'token_type_ids' in inputs:\n",
    "        del inputs['token_type_ids']\n",
    "    outputs = model.generate(    \n",
    "        **inputs,        \n",
    "        pad_token_id=tokenizer.pad_token_id,        \n",
    "        max_new_tokens=max_new_tokens,        \n",
    "        do_sample=True,        \n",
    "        top_k=40,       \n",
    "        top_p=0.8,        \n",
    "        temperature=0.8    \n",
    "    )   \n",
    "    generated_text = tokenizer.decode( \n",
    "        outputs[0],        \n",
    "        skip_special_tokens=True    \n",
    "    )    \n",
    "    # print(outputs)    \n",
    "    print(generated_text)\n",
    "inference(model, tokenizer, input_text = '类型#上衣*版型#宽松*颜色#粉红色*图案#字母*图案#文字*图案#线条*衣样式#卫衣*衣款式#不规则', max_new_tokens=320)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c17f059e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "yolo5",
   "language": "python",
   "name": "yolo5"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
